from crewai import Agent, Crew, Process, Task
from crewai.project import CrewBase, agent, crew, task, before_kickoff
from crewai.agents.agent_builder.base_agent import BaseAgent
from crewai_tools import FileReadTool
from typing import List
import yaml
import re
from datetime import datetime

from investment_team.llm_config import prediction_llm
from investment_team.tools.etf_data_tool import ETFDataTool
 
# 如果您想在团队启动之前或之后运行代码片段，
# 您可以使用 @before_kickoff 和 @after_kickoff 装饰器
# https://docs.crewai.com/concepts/crews#example-crew-class-with-decorators

@CrewBase
class InvestmentTeam():
    """InvestmentTeam crew"""

    agents: List[BaseAgent]
    tasks: List[Task]
    
    # 定义工具数组属性
    def __init__(self):
        # 初始化所有工具
        self.etf_data_tool = ETFDataTool()
        self.file_read_tool = FileReadTool()
        
        # 定义所有 agent 共享的工具配置
        self.shared_tools = [
            self.etf_data_tool,      # ETF数据工具
            self.file_read_tool,     # 文件读取工具
        ]
        
        # 加载 planning 配置
        try:
            with open('/home/zhangjianyong/project/crewai/investment_team/src/investment_team/config/planning.yaml', 'r', encoding='utf-8') as f:
                self.planning_config = yaml.safe_load(f)
        except Exception:
            self.planning_config = None
        
        # 初始化基金代码变量，用于文件名拼接
        self.fund_code = None
    
    def _extract_fund_code(self, fund_str: str) -> str:
        """
        从 fund 字符串中提取基金代码，并清理为适合文件名的格式。
        例如：'sz159570,港股通创新药ETF' -> 'sz159570'
        """
        if not fund_str:
            return 'default'
        
        # 提取基金代码（逗号前的部分）
        fund_code = fund_str.split(',')[0].strip()
        
        # 清理不适合文件名的字符，只保留字母、数字和下划线
        fund_code = re.sub(r'[^\w\-_]', '', fund_code)
        
        return fund_code if fund_code else 'default'
    
    def _parse_datetime(self, time_str: str) -> tuple:
        """
        从时间字符串中解析出日期和时间，返回格式化的字符串。
        返回: (date_str, time_part) 例如: ('20251102', '1511')
        """
        try:
            # 尝试解析时间字符串
            dt = datetime.fromisoformat(time_str.replace(' ', 'T', 1))
            date_str = dt.strftime('%Y%m%d')  # 格式: 20251102
            time_part = dt.strftime('%H%M')    # 格式: 1511
            return date_str, time_part
        except Exception:
            # 如果解析失败，使用当前时间
            dt = datetime.now()
            date_str = dt.strftime('%Y%m%d')
            time_part = dt.strftime('%H%M')
            return date_str, time_part
    
    def _fetch_etf_data(self, fund_code: str) -> str:
        """
        预获取ETF数据
        
        Args:
            fund_code: 基金代码
            
        Returns:
            str: 格式化后的ETF数据文本
        """
        try:
            # 调用etf_data_tool获取数据
            result = self.etf_data_tool._run([{"code": fund_code, "limit": 60}])
            return result
        except Exception as e:
            print(f"Warning: Failed to fetch ETF data: {e}")
            return f"ETF数据获取失败: {str(e)}"
    
    @before_kickoff
    def prepare_output_files(self, inputs: dict):
        """
        在团队启动前，根据 inputs 中的参数更新任务的 output_file 和描述中的占位符。
        同时预获取所有需要的数据并添加到任务上下文中。
        """
        # 从 inputs 中获取参数
        fund_str = inputs.get('fund', '')
        current_time = inputs.get('current_time', '')
        self.fund_code = self._extract_fund_code(fund_str)
        
        # 解析日期和时间
        date_str, time_str = self._parse_datetime(current_time)
        
        # 从 current_time 中提取年份
        try:
            dt = datetime.fromisoformat(current_time.replace(' ', 'T', 1))
            year = str(dt.year)
        except Exception:
            year = str(datetime.now().year)
        
        # 将 fund_code 和 year 添加到 inputs 中，以便在格式化任务描述时使用
        inputs_with_fund_code = {**inputs, 'fund_code': self.fund_code, 'year': year}
        
        # 预获取数据
        etf_data = self._fetch_etf_data(self.fund_code)
        
        # 格式化数据上下文
        data_context = f"""
【预获取数据上下文】

=== ETF基金数据 ===
{etf_data}

【重要提示】
以上数据已预获取并包含在上下文中，请直接使用这些数据进行分析，不需要调用任何工具。
"""
        
        # 更新所有任务的 output_file 和描述中的占位符
        # 注意：self.tasks 在 @before_kickoff 时应该已经通过 @task 装饰器初始化
        for task in self.tasks:
            # 替换任务描述和期望输出中的占位符
            try:
                if hasattr(task, 'description') and task.description:
                    task.description = task.description.format(**inputs_with_fund_code)
            except (AttributeError, KeyError) as e:
                # 如果格式化失败，记录错误但不中断执行
                print(f"Warning: Failed to format task description: {e}")
            
            try:
                if hasattr(task, 'expected_output') and task.expected_output:
                    task.expected_output = task.expected_output.format(**inputs_with_fund_code)
            except (AttributeError, KeyError) as e:
                # 如果格式化失败，记录错误但不中断执行
                print(f"Warning: Failed to format task expected_output: {e}")
            
            # 将预获取的数据添加到任务的描述中，这样agent可以直接看到数据
            if hasattr(task, 'description') and task.description:
                # 在任务描述的开头添加数据上下文
                task.description = data_context + "\n\n" + task.description
            
            # 通过任务的 agent 的角色名来识别对应的任务
            task_agent = getattr(task, 'agent', None)
            if task_agent:
                agent_role = getattr(task_agent, 'role', '') or ''
                
                # 合并后的任务直接输出到日期目录下
                # 匹配多种可能的 role 名称：投资预测分析师、预测分析师、投资分析专家
                if ('投资预测分析师' in agent_role or '预测分析师' in agent_role or 
                    '投资分析专家' in agent_role or 'ETF基金' in agent_role):
                    task.output_file = f'out/{date_str}/{self.fund_code}/{self.fund_code}_integrated_analysis_report_{date_str}_{time_str}.md'
    

    # 在此处了解更多关于 YAML 配置文件的信息：
    # 代理：https://docs.crewai.com/concepts/agents#yaml-configuration-recommended
    # 任务：https://docs.crewai.com/concepts/tasks#yaml-configuration-recommended
    
    # 如果您想为您的代理添加工具，您可以在此处了解更多信息：
    # https://docs.crewai.com/concepts/agents#agent-tools
    @agent
    def prediction(self) -> Agent:
        """
        角色: ETF基金投资预测分析师
        目标: 基于技术指标分析进行涨跌预测和生成买卖决策，提供具体的交易建议
        背景: 你是一个经验丰富的ETF基金投资预测分析师。
        你擅长基于常用技术指标（成交量、K线形态、KDJ、MACD、布林带、RSI、均线系统等）的分析结果，预测ETF基金的未来涨跌方向（上涨/下跌），并评估预测的置信度。
        你能够结合多个技术指标的综合评分，生成具体的买卖决策（买入/卖出/持有），并提供决策理由和时机建议。
        你以准确的技术指标解读能力和清晰的决策建议而闻名。
        """
        return Agent(
            config=self.agents_config['prediction'], # type: ignore[index]
            verbose=True, # 启用详细输出，便于调试和监控
            memory=False,   # 启用记忆功能，智能体可以记住之前的对话和上下文
            llm=prediction_llm,  # 使用预测分析师专用LLM
            allow_delegation=False,  # 专家agent专注于自己的专长
            tools=[]  # 移除所有工具，数据已预获取并在上下文中提供
        )

    # 要了解更多关于结构化任务输出、
    # 任务依赖关系和任务回调的信息，请查看文档：
    # https://docs.crewai.com/concepts/tasks#overview-of-a-task

    @task
    def integrated_analysis_task(self) -> Task:
        """
        整合的技术指标分析和预测决策任务
        所有数据已预获取并在上下文中提供，agent不需要调用任何工具
        """
        return Task(
            config=self.tasks_config['integrated_analysis_task'], # type: ignore[index]
            agent=self.prediction(),  # 使用预测分析师agent
            context=[],  # 数据将通过prepare_output_files方法添加到context中
            output_file='out/integrated_analysis_report.md'  # 默认输出文件路径，会在prepare_output_files中更新
        )

    @crew
    def crew(self) -> Crew:
        """创建投资团队"""
        # 要了解如何为您的团队添加知识源，请查看文档：
        # https://docs.crewai.com/concepts/knowledge#what-is-knowledge

        return Crew(
            agents=[self.prediction()],  # 只保留预测分析师agent
            tasks=self.tasks,  # 只包含整合后的任务
            process=Process.sequential,  # 顺序流程
            verbose=True
        )
